library(tidyverse)
library(readr)
library(maptools)
library(doMC)

ben_theme <- function(){
    theme_classic() +
        theme(panel.background = element_blank(),
              panel.grid.major = element_blank(),
              panel.grid.minor = element_blank(),
              axis.line = element_line(colour = "black"),
              panel.border = element_rect(colour = "black", fill = NA, size = 1),
              strip.background = element_blank(),
              legend.position = "bottom", legend.title = element_blank(),
              plot.title = element_text(hjust = 0.5))
}

generateLat <- function(nrow, ncol){
    myLat <- list()
    for(i in 1:nrow){
        for(j in 1:ncol){
            index <- ncol*(i-1) + j
            edges <- NULL
            if(i > 1){
                edges <- c(edges, index - ncol)
            }
            if(j > 1){
                edges <- c(edges, index -1)
            }
            if(j < ncol){
                edges <- c(edges, index +1)
            }			
            if(i < nrow){
                edges <- c(edges, index + ncol)
            }
            
            myLat[[index]] = edges - 1
        }
    }
    myLat
}

## Set parameters
params <- expand.grid(
    lat_size = seq(3, 15, by = 1),
    ndists = 2:3,
    alg = c("enumpart"),
    stringsAsFactors = FALSE
)

dir.create("../../data/runtime_test_lattice", showWarnings = FALSE)

## Get info on job
i <- as.numeric(Sys.getenv("SLURM_ARRAY_TASK_ID"))

## ---------
## Run tests
## ---------
## Get parameters
lat_size <- params$lat_size[i]
ndists <- params$ndists[i]
alg <- params$alg[i]

## Create lattice
adjlist <- generateLat(lat_size, lat_size)
adjlist <- lapply(adjlist, function(x){x+1})

## ## Runtime for redist.enumerate()
## if(alg == "redist"){
##     start <- Sys.time()
##     re_out <- redist.enumerate(nh_sub, ndists = ndists)
##     end <- Sys.time()
##     time <- difftime(end, start, units = "secs")
## }

## Runtime for enumpart
if(alg == "enumpart"){

    start <- Sys.time()

    ## Sink
    adjlist_map <- c()
    for(k in 1:length(adjlist)){
        sub <- adjlist[[k]]
        sub <- sub[sub > k]
        if(length(sub) > 0){
            for(l in 1:length(sub)){
                adjlist_map <- rbind(adjlist_map, c(k, sub[l]))
            }
        }
    }
    write_delim(
        data.frame(adjlist_map),
        path = paste0("../../data/runtime_test_lattice/adjlist_map_", i, ".dat"),
        col_names = FALSE
    )
    
    ## Order edges
    system(paste0("python ../ndscut.py <../../data/runtime_test_lattice/adjlist_map_", i, ".dat >../../data/runtime_test_lattice/adjlist_map_ordered_", i, ".dat"))

    ## Run enumpart
    system(paste0("../../enumpart_private/enumpart ../../data/runtime_test_lattice/adjlist_map_ordered_", i, ".dat -k ", ndists))

    end <- Sys.time()
    time <- difftime(end, start, units = "secs")
    
}

## Format and save
out <- data.frame(lat_size = lat_size, ndists = paste0(ndists, " Districts"),
                  alg = alg, runtime = as.numeric(time))
write_csv(out, path = paste0("../../data/runtime_test_lattice/runtime_test_", i, ".csv"))
